#!/usr/bin/env python3

import rclpy
import socket
import struct
import threading
from rclpy.node import Node
from std_msgs.msg import Float64MultiArray

class UDPBridge(Node):
    def __init__(self):
        super().__init__('udp_bridge')

        # 通过参数获取本地绑定和目标IP及端口
        self.declare_parameter("local_ip", "127.0.0.1")
        self.declare_parameter("local_port", 3000)
        self.declare_parameter("target_ip", "192.168.3.23")
        self.declare_parameter("target_port", 3001)
        self.local_ip = self.get_parameter("local_ip").value
        self.local_port = self.get_parameter("local_port").value
        self.target_ip = self.get_parameter("target_ip").value
        self.target_port = self.get_parameter("target_port").value

        # 初始化UDP socket，并绑定本地地址
        self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.udp_socket.bind((self.local_ip, self.local_port))
        self.udp_socket.settimeout(0.1)

        # ROS2话题
        self.send_data = self.create_publisher(Float64MultiArray, '/udp_send_data', 10) # 发送给其他节点的数据
        self.receive_data = self.create_subscription(
            Float64MultiArray, '/udp_receive_data', self.feedback_callback, 10) # 接收其他节点的数据并发送给上位机

        # UDP接收线程
        self.recv_thread = threading.Thread(target=self.udp_recv_loop)
        self.recv_thread.daemon = True
        self.recv_thread.start()

        # 定时UDP发送，20Hz
        self.feedback_data = None
        self.timer = self.create_timer(0.05, self.udp_send_loop)

        # 结构体格式
        self.RECV_STRUCT_FORMAT = ">IIIIIIfffffffffI"
        self.RECV_STRUCT_SIZE = struct.calcsize(self.RECV_STRUCT_FORMAT)
        self.SEND_STRUCT_FORMAT = ">IIIIffffffffffffI"
        self.SEND_STRUCT_SIZE = struct.calcsize(self.SEND_STRUCT_FORMAT)

    def udp_recv_loop(self):
        while rclpy.ok():
            try:
                data, addr = self.udp_socket.recvfrom(1024)
                print(f"Received data from {addr}: {data}")
                # 只处理来自目标IP和端口的数据
                if addr[0] != self.target_ip or addr[1] != self.target_port:
                    continue
                self.get_logger().info(f"Received data from {addr}: {data}")
                unpacked = struct.unpack(self.RECV_STRUCT_FORMAT, data)
                # 验证头尾
                if unpacked[0] == 0x55555555 and unpacked[-1] == 0xaaaaaaaa:
                    msg = Float64MultiArray()
                    msg.data = list(unpacked[1:-1])  # 去掉头尾
                    self.send_data.publish(msg)
            except socket.timeout:
                continue
            except Exception as e:
                self.get_logger().warn(f"UDP recv error: {e}")

    def feedback_callback(self, msg):
        # 保存最新的反馈数据
        self.feedback_data = msg.data

    def udp_send_loop(self):
        if self.feedback_data is None or len(self.feedback_data) < 14:
            return
        head = 0x66666666
        tail = 0x88888888
        try:
            data = struct.pack(
                self.SEND_STRUCT_FORMAT,
                head,
                int(self.feedback_data[0]),  # count
                int(self.feedback_data[1]),  # status
                0,  # 保留
                float(self.feedback_data[2]),  # x
                float(self.feedback_data[3]),  # y
                float(self.feedback_data[4]),  # z
                float(self.feedback_data[5]),  # roll
                float(self.feedback_data[6]),  # yaw
                float(self.feedback_data[7]),  # pitch
                float(self.feedback_data[8]),  # l1
                float(self.feedback_data[9]),  # l2
                float(self.feedback_data[10]), # l3
                float(self.feedback_data[11]), # l4
                float(self.feedback_data[12]), # l5
                float(self.feedback_data[13]), # l6
                tail
            )
            self.udp_socket.sendto(data, (self.target_ip, self.target_port))
        except Exception as e:
            self.get_logger().warn(f"UDP send error: {e}")

def main(args=None):
    rclpy.init(args=args)
    node = UDPBridge()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()